/*
题目: 将矩阵按对角线排序
矩阵对角线 是一条从矩阵最上面行或者最左侧列中的某个元素开始的对角线，沿右下方向一直到矩阵末尾的元素。例如，矩阵 mat 有 6 行 3 列，从 mat[2][0] 开始的 矩阵对角线 将会经过 mat[2][0]、mat[3][1] 和 mat[4][2] 。

给你一个 m * n 的整数矩阵 mat ，请你将同一条 矩阵对角线 上的元素按升序排序后，返回排好序的矩阵。

https://leetcode.cn/problems/sort-the-matrix-diagonally
 */
public class DiagonalSort {
    // 计数排序优化
    public int[][] diagonalSortBetter(int[][] mat) {
        // 数据较小, 计数排序
        int row = mat.length;
        int col = mat[0].length;

        for (int i = 0; i < col; i++) {
            int[] count = new int[101];
            for (int x = 0, y = i; x < row && y < col; ) {
                count[mat[x ++][y ++]] ++;
            }

            int x = 0, y = i;
            for (int pos = 1; pos < 101; ) {
                while (count[pos]-- > 0) {
                    mat[x ++][y ++] = pos;
                }
                pos ++;
            }
        }

        for (int i = 1; i < row; i++) {
            int[] count = new int[101];
            for (int x = i, y = 0; x < row && y < col; ) {
                count[mat[x ++][y ++]] ++;
            }

            int x = i, y = 0;
            for (int pos = 1; pos < 101; ) {
                while (count[pos]-- > 0) {
                    mat[x ++][y ++] = pos;
                }
                pos ++;
            }
        }

        return mat;
    }

    // 基础版: 直接取出来然后排序
    public int[][] diagonalSort(int[][] mat) {
        int row = mat.length;
        int col = mat[0].length;

        for (int i = 0; i < col; i ++) {
            ArrayList<Integer> cur = new ArrayList<> ();
            for (int x = 0, y = i; x < row && y < col; ) {
                cur.add(mat[x++][y++]);
            }
            Collections.sort(cur);
            for (int x = 0, y = i, pos = 0; x < row && y < col; ) {
                mat[x++][y++] = cur.get(pos ++);
            }
        }

        for (int i = 0; i < row; i ++) {
            ArrayList<Integer> cur = new ArrayList<> ();
            for (int x = i, y = 0; x < row && y < col; ) {
                cur.add(mat[x ++][y ++]);
            }
            Collections.sort(cur);
            for (int x = i, y = 0, pos = 0; x < row && y < col; ) {
                mat[x ++][y ++] = cur.get(pos ++);
            }
        }

        return mat;
    }
}
